import json
import time
import logging
import paho.mqtt.client as mqtt

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


class SimpleMQTT:
    def __init__(self, host="127.0.0.1", port=1883, username=None, password=None):
        self.host, self.port = host, port
        self.username, self.password = username, password
        self.client = mqtt.Client()

        if username and password:
            self.client.username_pw_set(username, password)

        self.connected = False
        self.client.on_connect = self._on_connected
        self.client.on_disconnect = self._on_disconnected
        self.client.on_message = self._on_message
        self.subscriptions = {}
        self._connect()

    def _on_connected(self, client, userdata, flags, rc):
        if rc == 0:
            self.connected = True
            logger.info(f"[MQTT] Connected to {self.host}:{self.port}")
            for topic in self.subscriptions.keys():
                client.subscribe(topic)
                logger.info(f"[MQTT] Re-subscribed to topic: {topic}")
        else:
            logger.warning(f"[MQTT] Connection failed, code={rc}")

    def _on_disconnected(self, client, userdata, rc):
        self.connected = False
        logger.info(f"[MQTT] Disconnected (rc={rc})")

    def _on_message(self, client, userdata, msg):
        topic = msg.topic
        payload = msg.payload.decode("utf-8", errors="ignore")
        logger.info(f"[MQTT] Message received: {topic} => {payload}")
        cb = self.subscriptions.get(topic)
        if cb:
            try:
                cb(topic, payload)
            except Exception as e:
                logger.exception(f"[MQTT] Error in callback for {topic}: {e}")

    def _connect(self):
        try:
            self.client.connect(self.host, self.port, keepalive=60)
            self.client.loop_start()
        except Exception as e:
            logger.error(f"[MQTT] Connect failed: {e}")
            self.connected = False

    def wait_connected(self, timeout=5):
        """等待连接成功（阻塞）"""
        start = time.time()
        while not self.connected and time.time() - start < timeout:
            time.sleep(0.1)
        if not self.connected:
            logger.warning(f"[MQTT] Wait connected timeout ({timeout}s)")
        return self.connected

    def publish(self, payload, topic=None, qos=1, retain=False):
        if not self.connected:
            logger.warning("[MQTT] Not connected, retrying...")
            self._connect()
            time.sleep(1)

        payload_str = (
            json.dumps(payload, ensure_ascii=False)
            if isinstance(payload, (dict, list))
            else str(payload)
        )
        if self.connected and topic:
            self.client.publish(topic, payload_str, qos=qos, retain=retain)
            logger.info(f"[MQTT] Published to {topic}: {payload_str}")

    def subscribe(self, topic, callback=None, qos=1):
        """安全订阅：连接成功后再订阅"""
        if not self.connected:
            logger.info(f"[MQTT] Waiting for connection before subscribing {topic}...")
            self.wait_connected(timeout=5)
        if self.connected:
            self.client.subscribe(topic, qos=qos)
            self.subscriptions[topic] = callback
            logger.info(f"[MQTT] Subscribed to topic: {topic}")
        else:
            logger.warning(f"[MQTT] Subscribe failed, still not connected.")

    def stop(self):
        self.client.loop_stop()
        self.client.disconnect()
        logger.info("[MQTT] Client stopped.")

def on_message_callback(topic, payload):
    logger.info(f"[回调] 收到消息：{topic} => {payload}")

if __name__ == "__main__":
    mqtt_client = SimpleMQTT(host="192.168.6.7", port=2771,username='edge8563',password='XnZtSd1WlvF7Qt2E')

    recv_topic = "/iios/coco/mqtt-send/session/04f5c05f-db68-4ec4-9368-ed0ef9c60dda/msg"

    send_topic = "/iios/coco/mqtt-send/session/86e92aaa-1e6c-4eb3-823e-8a018e88a288/msg"

    # 订阅振动传感器采集的消息
    mqtt_client.wait_connected(timeout=5)
    mqtt_client.subscribe(recv_topic, on_message_callback)
    # 这里是推理代码
    pass
    # 将推理的结果发送给平台。data的key可以自己定义
    try:
       while True:
           data = {"device": "demo01", "value": round(time.time() % 100, 2)}
           mqtt_client.publish(data, topic=send_topic)
           time.sleep(1)
    except KeyboardInterrupt:
       mqtt_client.stop()
       logger.info("[MQTT] 程序结束。")
